#include <istream>
#include <vector>
#include <list>
#include <string>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <map>
#include <list>
#include <pthread.h>
#include <boost/shared_array.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/format.hpp>
#include <stdlib.h>

#include "components/UPGMA_nearest_neighbor.h"
#include "components/bcpp_tree.h"
#include "components/portable_timer.h"
#include "components/perlutil.h"

using boost::format;
using std::vector;
using std::ostringstream;
using boost::shared_array;
using boost::shared_ptr;

using std::max;
using std::min;
using std::string;
using std::fill;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::ifstream;
using std::map;
using std::pair;
using std::istream;
using std::list;

struct t_ks_pairs
{
	double		ks;
	unsigned gene1;
	unsigned gene2;
	t_ks_pairs(double d, unsigned g1, unsigned g2): ks(d), gene1(g1), gene2(g2){}
};


void allocate_matrix(vector<shared_array<double> >& matrix, unsigned size)
{
	// get data and names
	for (unsigned j = 0; j < size; ++j)
	{
		matrix.push_back(shared_array<double>(new double [size]));
		fill(matrix.back().get(), matrix.back().get() + size, 10000.0);
	}
}

string parse_cluster(	istream& infile, vector<t_ks_pairs>& ks_pairs, vector<string>& names,
						map<string, unsigned>& index_to_names )
{
	string line;
	string cluster_name;
	while (getline(infile, line))
	{
		if (!line.length())
		{
			// ignore consecutive blank lines
			if (!names.size())
				continue;

			// blank lines separate clusters
			else
				break;
		}
		if (line[0] == '>')
		{
			cluster_name = line.substr(1);;
			continue;
		}

		vector<string> data;
		perl_split (line, "\t", data, false, 3);



		pair<map<string, unsigned>::iterator, bool> pr;
		typedef pair<string, unsigned> t_gene_index_pair;

		//
		// first gene in pair
		//
		// save name if new, else retrieve existing index to name
		unsigned gene1_index = names.size();
		pr = index_to_names.insert(t_gene_index_pair(data[0],  gene1_index));
		if (pr.second)
			names.push_back(data[0]);
		else
			gene1_index = (pr.first)->second;

		//
		// second gene in pair
		//
		// save name if new, else retrieve existing index to name
		unsigned gene2_index = names.size();
		pr = index_to_names.insert(t_gene_index_pair(data[1],  gene2_index));
		if (pr.second)
			names.push_back(data[1]);
		else
			gene2_index = (pr.first)->second;

		ks_pairs.push_back(t_ks_pairs(atof(data[2].c_str()), gene1_index, gene2_index));
	}
	return cluster_name;
}


void create_main_matrix(vector<shared_array<double> >& matrix,
						vector<t_ks_pairs>& ks_pairs,
						unsigned cnt_leaves)
{
	// get data and names
	allocate_matrix(matrix, cnt_leaves);
	for (unsigned i = 0; i < ks_pairs.size(); ++i)
	{
		assert(ks_pairs[i].gene1 < cnt_leaves);
		assert(ks_pairs[i].gene2 < cnt_leaves);
		matrix[ks_pairs[i].gene1][ks_pairs[i].gene2]=
		matrix[ks_pairs[i].gene2][ks_pairs[i].gene1]= ks_pairs[i].ks;
	}
}


void make_clusters(vector<t_ks_pairs>& ks_pairs, unsigned cnt_items,
					double cutoff,
					vector<vector<unsigned> >& cluster_set)
{
	//// NB cluster_ids start from 1: zero = unassigned cluster
	vector<unsigned>			gene_cluster_ids (cnt_items, 0);
	vector<list<unsigned> >		gene_clusters(1);

	for (unsigned i = 0; i < ks_pairs.size(); ++i)
	{
		if (ks_pairs[i].ks > cutoff)
			continue;
		if (ks_pairs[i].ks < 0.0)
			continue;
		cerr << ks_pairs[i].ks <<
				"\t" << ks_pairs[i].gene1<<
			"\t" << ks_pairs[i].gene2<<
			"\n";
		unsigned gene1 = ks_pairs[i].gene1;
		unsigned gene2 = ks_pairs[i].gene2;
		unsigned cluster_id1 = gene_cluster_ids[gene1];
		unsigned cluster_id2 = gene_cluster_ids[gene2];

		// add new cluster if neither is yet clustered
		if (!(cluster_id1 + cluster_id2))
		{
			// cluster index for these two genes
			gene_cluster_ids[gene1] =
			gene_cluster_ids[gene2] = gene_clusters.size();

			// these two genes make up a new cluster
			gene_clusters.push_back(list<unsigned>());
			gene_clusters.back().push_back(gene1);
			gene_clusters.back().push_back(gene2);
		}

		// link cluster 2 to cluster 1 if both already have clusters
		else if (cluster_id1 && cluster_id2)
		{

			// assign cluster_id1 to all of cluster_id2 genes;
			list<unsigned>::iterator genes_iter = gene_clusters[cluster_id2].begin();
			list<unsigned>::iterator genes_end = gene_clusters[cluster_id2].end();
			for (; genes_iter != genes_end; ++genes_iter)
				gene_cluster_ids[*genes_iter] = cluster_id1;

			// merge cluster 2 into cluster 1
			gene_clusters[cluster_id1].splice(	gene_clusters[cluster_id1].end(),
												gene_clusters[cluster_id2]);
		}

		// add gene 2 to cluster 1 if gene 2 is unclustered
		else if (cluster_id1)
		{
			gene_clusters[cluster_id1].push_back(gene2);
			gene_cluster_ids[gene2] = cluster_id1;
		}

		// add gene 1 to cluster 2 if gene 1 is unclustered
		else
		{
			gene_clusters[cluster_id2].push_back(gene1);
			gene_cluster_ids[gene1] = cluster_id2;
		}
	}

	//
	// copy non-empty clusters to output set
	//
	// remember: cluster 0 is a dummy
	for (unsigned i = 1; i < gene_clusters.size(); ++i)
	{
		if (gene_clusters[i].empty())
			continue;
		cluster_set.push_back(vector<unsigned>());
		cluster_set.back().insert(cluster_set.back().end(),
									gene_clusters[i].begin(),
									gene_clusters[i].end());

		list<unsigned>::iterator j = gene_clusters[i].begin();
		for (; j != gene_clusters[i].end(); j++)
		{
			cerr << *j << "!\n";
		}
		cerr << "!#\n";

	}
}


void cluster_matrix_and_names(	vector<unsigned>& 						genes,
								vector<shared_array<double> >&			matrix,
								vector<string>&							names,
								vector<vector<shared_array<double> > >& vec_matrices,
								vector<vector<string> >& 				vec_names)

{
	// add names
	vec_names.push_back(vector<string>());
	vector<string>& new_names = vec_names.back();
	for (unsigned i = 0; i < genes.size(); ++i)
		new_names.push_back(names[genes[i]]);

	// add matrix
	vec_matrices.push_back(vector<shared_array<double> >());
	vector<shared_array<double> >& new_matrix = vec_matrices.back();
	allocate_matrix(new_matrix, genes.size());
	for (unsigned i = 0; i < genes.size(); ++i)
		for (unsigned j = 0; j < i; ++j)
			new_matrix[i][j] =
			new_matrix[j][i] = matrix[genes[i]][genes[j]];
}

unsigned getdata(istream& infile,
				vector<vector<shared_array<double> > >& vec_matrices,
				vector<vector<string> >& vec_names,
				double cutoff, unsigned& cluster_size,
				string& cluster_name)
{
	vector<t_ks_pairs>				ks_pairs;
	vector<string>					names;
	map<string, unsigned>			index_to_names;
	vector<shared_array<double> >	matrix;

	cluster_name = parse_cluster(infile, ks_pairs, names, index_to_names);
	create_main_matrix(matrix, ks_pairs, names.size());

	vector<vector<unsigned> > clusters;
	make_clusters(ks_pairs, names.size(), cutoff, clusters);

	for (unsigned i = 0; i < clusters.size(); i++)
	{
		for (unsigned j = 0; j < clusters[i].size(); ++j)
		{
			cerr << i << "--" << names [clusters[i][j]] << "\n";
		}
		cluster_matrix_and_names(clusters[i], matrix, names, vec_matrices, vec_names);
	}
	cluster_size = names.size();
	return ks_pairs.size();
}  /* getdata */



//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		str_tree_maker

//________________________________________________________________________________________
//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//________________________________________________________________________________________

//	print_orthologs()

//________________________________________________________________________________________
bool open_filename(ofstream& ofs, const string& file_name)
{
	ofs.open(file_name.c_str());
	if (ofs)
		return true;
	cerr << "Error\nCould not open " << file_name << "\n";
	return false;
}


struct t_outputstreams
{
	t_outputstreams(const string& f): cluster_count(0), file_prefix(f)
	{
		open();
		ofs_clusters << "Cluster ID\tCluster size\tNumber of duplicated genes\n";
		ofs_all_nodes << "Branch length /\tFrom root /\tFrom leaf\n" << string(80, '-') << "\n";

		ofs_full_tree << "Branch length /\tFrom root /\tFrom leaf\n" << string(80, '-') << "\n";
	}
	ofstream ofs_genes;
	ofstream ofs_full_tree;
	ofstream ofs_newhampshire_tree;
	ofstream ofs_all_nodes;
	ofstream ofs_clusters;
	unsigned cluster_count;
	string	 file_prefix;
	bool open();
	bool analyse_tree(bcpp_tree::t_tree& tree, const string& cluster_name,
						unsigned subcluster, unsigned cluster_size);
};


bool t_outputstreams::analyse_tree(bcpp_tree::t_tree& tree, const string& cluster_name,
									unsigned subcluster, unsigned cluster_size)
{
//	cerr << "\n" << string(80, '*') << "\n\n";
	cluster_count++;
	unsigned sub_cluster_size = tree.get_children_count();
	if (sub_cluster_size < 2)
		return true;
	double max_ks = tree.get_leaf(0).get_len_to_root();
//	if (max_ks > 2500.0)
//		return true;

	// dump all nodes
	{
		ofs_all_nodes << "Cluster id = " << cluster_name << "." << subcluster << "\n";
		bcpp_tree::do_dump_nodes dump_nodes(ofs_all_nodes);
		tree.for_each(dump_nodes);
		ofs_all_nodes << "\n";
	}

	// print full tree
	{

		ofs_full_tree << "Cluster id = " << cluster_name << "." << subcluster << "\n";
		bcpp_tree::do_print_tree_full print_tree(ofs_full_tree);
		tree.for_each(print_tree);
		ofs_full_tree << "\n";
	}

	// new hampshire tree
	{
		ofs_newhampshire_tree << "Cluster id = " << cluster_name << "." << subcluster << "\n";
		bcpp_tree::do_newhampshire_format_str nh_str(ofs_newhampshire_tree);
		tree.for_each(nh_str);
		ofs_newhampshire_tree << "\n";
	}

	// print gene sets
	{
//		ofs_genes << "Cluster id = " << cluster_name << "." << subcluster << "\n";
		for (unsigned i = sub_cluster_size; i !=0;)
		{
			--i;
			ofs_genes	<< tree.get_leaf(i).get_name() << "\t"
						<< cluster_name 			 << "\t"
						<< subcluster                << "\t"
						<< sub_cluster_size          << "\t"
						<< max_ks                    << "\n";
		}
	}

	// print gene sets
	{
		ofs_clusters    << cluster_name						<< "\t"
						<< cluster_size						<< "\t"
						<< sub_cluster_size 				<< "\n";
	}
	return true;

}



bool t_outputstreams::open()
{
	return(
			open_filename(ofs_genes,	 		file_prefix + "genes.list"			)	&&
			open_filename(ofs_clusters, 		file_prefix + "clusters.list"		)	&&
			open_filename(ofs_full_tree, 		file_prefix + "full_topology.tree"	)	&&
			open_filename(ofs_newhampshire_tree,file_prefix + "new_hampshire.tree"	)	&&
			open_filename(ofs_all_nodes, 		file_prefix + "all_nodes.list"		)
		);
}

struct missing_pairs
{
	unsigned cluster_num;
	unsigned actual_pairs;
	unsigned expected_pairs;
	missing_pairs(unsigned c, unsigned a, unsigned e):
		cluster_num(c), actual_pairs(a), expected_pairs(e){}
};


int main (int argc, char *argv[])
{
	std::ios::sync_with_stdio(false);


	if (argc < 3 || argc > 4)
	{
		std::cerr << "Usage:\n\tparology_tree <cutoff> <file_prefix> [infile | <infile]\n"
					 "\twhere cutoff is a floating point threshold for clustering genes.\n\n";
		return 1;
	}
	double cutoff	 	= atof(argv[1]);
	string file_prefix	= argv[2];
	istream* datafile	= &cin;
	ifstream infile;
	if (argc >= 4)
	{
		string infile_name = argv[3];
		infile.open(infile_name.c_str());
		if (!infile)
		{
			std::cerr << "Error opening [" << infile_name <<"] \n";
			return 1;
		}
		datafile = &infile;
	}


	unsigned cnt_cluster = 0;
	vector<missing_pairs> problems;
	t_outputstreams strms(file_prefix);

	{
		t_progress_indicator dots(cerr, 50);
		while (datafile->good())
		{
			++cnt_cluster;
			vector<vector<shared_array<double> > >	vec_matrices;
			vector<vector<string> > 				vec_names;
			string									cluster_name;
			unsigned								cluster_size;
			unsigned cnt_pairs = getdata(*datafile, vec_matrices, vec_names, cutoff,
											cluster_size, cluster_name);
			assert(vec_matrices.size() == vec_names.size());
			if (!cnt_pairs)
				break;
			unsigned expected_pairs = cluster_size * (cluster_size - 1) / 2;
			if (cnt_pairs != expected_pairs)
				problems.push_back(missing_pairs(cnt_cluster, cnt_pairs, expected_pairs));

			for (unsigned i = 0; i < vec_names.size(); i++)
			{
				bcpp_tree::t_tree new_tree;
				upgma_tree::upgma_make_tree(new_tree, vec_names[i], vec_matrices[i], true, dots);

				if (!strms.analyse_tree(new_tree, cluster_name, i, cluster_size))
					return 0;
			}
		}
	}
	if (problems.size())
	{
		cerr << problems.size() << "\tmatrixes had missing pairs:\nCheck " << file_prefix
					<< "missing_pairs.error\n";
		ofstream ofs_missing_pairs;
		if (open_filename(ofs_missing_pairs, file_prefix + "missing_pairs.error"))
		{
			for (unsigned i = 0; i < problems.size(); ++i)
			{
				ofs_missing_pairs
					<< format("\tCluster #%|1$+4|: Only %|2$+4| out of %|3$+4| pairs provided.\n")
						% problems[i].cluster_num
						% problems[i].actual_pairs
						% problems[i].expected_pairs;
			}
		}

	}
	cerr << "\nAll finished!!\n\n";

	return 0;
}


